import glob
import os
import os.path as osp
import torch
import numpy as np
from torch_geometric.io import read_txt_array
from torch_geometric.data import Data


def read_file(folder, prefix, name, dtype=None):
    path = osp.join(folder, f'{prefix}_{name}.txt')
    return read_txt_array(path, sep=',', dtype=dtype)


def cat(seq):
    seq = [item for item in seq if item is not None]
    seq = [item for item in seq if item.numel() > 0]
    seq = [item.unsqueeze(-1) if item.dim() == 1 else item for item in seq]
    return torch.cat(seq, dim=-1) if len(seq) > 0 else None


def split(data, batch):
    node_slice = torch.cumsum(torch.from_numpy(np.bincount(batch)), 0)
    node_slice = torch.cat([torch.tensor([0]), node_slice])

    row, _ = data.edge_index
    edge_slice = torch.cumsum(torch.from_numpy(np.bincount(batch[row])), 0)
    edge_slice = torch.cat([torch.tensor([0]), edge_slice])

    # Edge indices should start at zero for every graph.
    data.edge_index -= node_slice[batch[row]].unsqueeze(0)

    slices = {'edge_index': edge_slice}
    if data.x is not None:
        slices['x'] = node_slice
    else:
        # Imitate `collate` functionality:
        data._num_nodes = torch.bincount(batch).tolist()
        data.num_nodes = batch.numel()

    if data.y is not None:
        if data.y.size(0) == batch.size(0):
            slices['y'] = node_slice
        else:
            slices['y'] = torch.arange(0, batch[-1] + 2, dtype=torch.long)

    return data, slices


def read_data(folder, prefix):
    files = glob.glob(osp.join(folder, f'{prefix}_*.txt'))
    names = [f.split(os.sep)[-1][len(prefix) + 1:-4] for f in files]

    edge_index = read_file(folder, prefix, 'A', torch.long).t() - 1
    batch = read_file(folder, prefix, 'graph_indicator', torch.long) - 1
    node_labels = torch.empty((batch.size(0), 0))
    node_attributes = torch.empty((batch.size(0), 0))

    if 'node_attributes' in names:
        node_attributes = read_file(folder, prefix, 'node_attributes')
        if node_attributes.dim() == 1:
            node_attributes = node_attributes.unsqueeze(-1)

    x = cat([node_attributes, node_labels])
    edge_attr = None

    y = None
    if 'graph_attributes' in names:  # Regression problem.
        y = read_file(folder, prefix, 'graph_attributes')
    elif 'graph_labels' in names:  # Classification problem.
        y = read_file(folder, prefix, 'graph_labels', torch.long)
        _, y = y.unique(sorted=True, return_inverse=True)

    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)
    data, slices = split(data, batch)

    return data, slices


if __name__ == '__main__':
    # lines_to_extract = 4580731
    # output_file = "output.txt"
    #
    # with open(r"D:\PY_test\Causal_Open_set\data\NPS\raw\NPS_A.txt", "r") as input_file, open(output_file, "w") as output_file:
    #     for i, line in enumerate(input_file):
    #         if i >= lines_to_extract:
    #             break
    #         output_file.write(line)

    data, slices = read_data(r"F:\论文\TII\Causal_Open_set\data\NPS\raw", "NPS")
    print(data)
